package com.pig4cloud.pigx.knowledge.service.impl;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.pig4cloud.pigx.common.core.constant.enums.YesNoEnum;
import com.pig4cloud.pigx.common.core.util.R;
import com.pig4cloud.pigx.knowledge.entity.AiDatasetEntity;
import com.pig4cloud.pigx.knowledge.mapper.AiDatasetMapper;
import com.pig4cloud.pigx.knowledge.mapper.AiModelMapper;
import com.pig4cloud.pigx.knowledge.service.AiDatasetService;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import dev.langchain4j.model.embedding.DimensionAwareEmbeddingModel;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.Objects;

/**
 * 知识库
 *
 * @author pig
 * @date 2024-03-14 13:39:21
 */
@Service
@RequiredArgsConstructor
public class AiDatasetServiceImpl extends ServiceImpl<AiDatasetMapper, AiDatasetEntity> implements AiDatasetService {

	private final QdrantClient qdrantClient;

	private final AiModelMapper modelMapper;

	private final ModelProvider modelProvider;

	/**
	 * “获取 AI 数据集”页面
	 * @param page 页
	 * @return {@link Page }<{@link AiDatasetEntity }>
	 */
	@Override
	public Page<AiDatasetEntity> getAiDatasetPage(Page<AiDatasetEntity> page) {
		return baseMapper.getAiDatasetPage(page);
	}

	/**
	 * 根据ID查询知识库
	 * @param datasetId 知识库ID
	 * @return R
	 */
	@Override
	public R getDatasetById(Long datasetId) {
		if (datasetId == null) {
			return R.failed("目标知识库ID不能为空");
		}

		AiDatasetEntity dataset = baseMapper.selectById(datasetId);
		if (Objects.isNull(dataset) || YesNoEnum.NO.getCode().equals(dataset.getPublicFlag())) {
			return R.failed(StrUtil.format("私有知识库，无法查看"));
		}
		// 移除敏感字段
		dataset.setPublicPassword(null);
		return R.ok(dataset);
	}

	/**
	 * 创建知识库
	 * @param aiDataset 知识库
	 * @return boolean
	 */
	@SneakyThrows
	@Override
	@Transactional(rollbackFor = Exception.class)
	public R saveDataset(AiDatasetEntity aiDataset) {
		this.save(aiDataset);
		// 创建集合,使用知识库的ID作为名称
		DimensionAwareEmbeddingModel embeddingModel = modelProvider.getEmbeddingModel(aiDataset.getEmbeddingModel());
		var vectorParams = Collections.VectorParams.newBuilder()
			.setDistance(Collections.Distance.Cosine)
			.setSize(embeddingModel.dimension())
			.build();
		this.qdrantClient.createCollectionAsync(aiDataset.getId().toString(), vectorParams).get();
		aiDataset.setCollectionName(aiDataset.getId().toString());
		baseMapper.updateById(aiDataset);
		return R.ok();
	}

	/**
	 * 集合是否存在
	 * @param collectionName 集合名称
	 * @return boolean
	 */
	private boolean isCollectionExists(String collectionName) {
		try {
			return this.qdrantClient.listCollectionsAsync().get().stream().anyMatch(c -> c.equals(collectionName));
		}
		catch (Exception e) {
			throw new RuntimeException(e);
		}
	}

	/**
	 * list 向量数据库名称列表
	 * @return {@link R }
	 */
	@Override
	@SneakyThrows
	public R listCollection() {
		return R.ok(qdrantClient.listCollectionsAsync().get());
	}

}
